package net.prasenjit.poc.httpclient;

import org.apache.http.conn.ssl.TrustStrategy;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.util.Assert;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Paths;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;

/**
 * Created by PPurohit on 12/22/2014.
 */
public class CustomSSLTrustStrategy implements TrustStrategy, InitializingBean {

    private KeyStore clientKeystore;
    private KeyStore defaultStore;
    private String clientKeystorePassword;

    @Override
    public boolean isTrusted(X509Certificate[] chain, String authType) throws CertificateException {
        try {
            boolean aliasFound = false;
            X509Certificate root = chain[chain.length - 1];
            if (defaultStore != null) {
                aliasFound |= defaultStore.getCertificateAlias(root) != null;
            }
            if (!aliasFound && null != clientKeystore) {
                aliasFound |= clientKeystore.getCertificateAlias(root) != null;
            }
            if (!aliasFound) return aliasFound;
            for (X509Certificate certificate : chain) {
                certificate.checkValidity();
            }
            return aliasFound;
        } catch (KeyStoreException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        Assert.notNull(clientKeystorePassword);
        File certStore = getDefaultKeyStoreFile();
        defaultStore = loadKeyStore(certStore, getDefaultKeyStoreType(), getKeyStorePassword());
        certStore = getClientKeyStoreFile();
        clientKeystore = loadKeyStore(certStore, KeyStore.getDefaultType(), clientKeystorePassword.toCharArray());
    }

    private File getClientKeyStoreFile() {
        String gcomPropDir = System.getProperty("gcomproperties.dir");
        return Paths.get(gcomPropDir, "clientTrust.jks").toFile();
    }

    private File getDefaultKeyStoreFile() {
        File certStore = null;
        String trustStoreProperty = System.getProperty("javax.net.ssl.trustStore");
        if (null == trustStoreProperty) {
            String javaHome = System.getProperty("java.home");
            certStore = Paths.get(javaHome, "lib/security/jssecacerts").toFile();
            if (!certStore.exists()) {
                certStore = Paths.get(javaHome, "lib/security/cacerts").toFile();
            }
        } else {
            certStore = Paths.get(trustStoreProperty).toFile();
        }
        return certStore;
    }

    private KeyStore loadKeyStore(File certStore, String keyStoreType, char[] keyStorePassword) {
        KeyStore keyStore = null;
        if (!certStore.exists()) return null;
        try (InputStream is = new FileInputStream(certStore);) {
            keyStore = KeyStore.getInstance(keyStoreType);
            keyStore.load(is, keyStorePassword);
            return keyStore;
        } catch (KeyStoreException | IOException | CertificateException | NoSuchAlgorithmException e) {
            throw new RuntimeException(e);
        }
    }

    private String getDefaultKeyStoreType() {
        String type = System.getProperty("javax.net.ssl.trustStoreType");
        if (type != null) {
            return type;
        }
        return KeyStore.getDefaultType();
    }

    private char[] getKeyStorePassword() {
        String password = System.getProperty("javax.net.ssl.trustStorePassword");
        if (null == password) {
            password = "changeit";
        }
        return password.toCharArray();
    }

    public void setClientKeystorePassword(String clientKeystorePassword) {
        this.clientKeystorePassword = clientKeystorePassword;
    }
}
